package rootkit

import (
	"log"
	"fmt"
	"net"
	"syscall"
	"strings"
	"strconv"
	"time"
	"cloud_worm/util/uexec"
)

// 对于机器上开放端口port的检查，主要表现为：
//		1. 利用socket的bind方式，检测open的port信息；
//  	2. 然后基于open的信息，利用netstat查看open的端口；
// 	结论：
//		1. 如果open的port和netstat不匹配，说明存在异常；

// 该rootkit的检查实际是确认 netstat 没问题

const (
	UDPProto = iota
	TCPProto
)

type PortCheck struct {
	NetStatTCP 		map[int]struct{}
	NetStatUDP 		map[int]struct{}
}

func (p *PortCheck) runNetStatAll() error {
	tcpCmd := "netstat -an | grep \"^tcp\""
	msg, ret := uexec.ExecUnixCmd(tcpCmd, 60)
	if ret != 0 {
		return fmt.Errorf(msg)
	} else if strings.Contains(msg, "exit status") {
		return fmt.Errorf(msg)
	}

	l := strings.Split(msg, "\n")
	for _, item := range l {
		item := strings.TrimSpace(item)
		if item == "" {
			continue
		}

		tmps := make([]string, 0, 6)
		splits := strings.Split(item, " ")
		for _, i := range splits {
			i := strings.TrimSpace(i)
			if i == "" {
				continue
			}
			tmps = append(tmps, i)
		}
		if len(tmps) == 6 {
			index := strings.LastIndex(tmps[3], ".")
			if index == -1 {
				continue
			}
			if port, err := strconv.Atoi(tmps[3][index+1:]); err == nil {
				p.NetStatTCP[port] = struct{}{}
			}
		}
	}

	udpCmd := "netstat -an | grep \"^udp\""
	msg, ret = uexec.ExecUnixCmd(udpCmd, 60)
	if ret != 0 {
		return fmt.Errorf(msg)
	} else if strings.Contains(msg, "exit status") {
		return fmt.Errorf(msg)
	}

	l = strings.Split(msg, "\n")
	for _, item := range l {
		item := strings.TrimSpace(item)
		if item == "" {
			continue
		}

		tmps := make([]string, 0, 6)
		splits := strings.Split(item, " ")
		for _, i := range splits {
			i := strings.TrimSpace(i)
			if i == "" {
				continue
			}
			tmps = append(tmps, i)
		}
		if len(tmps) == 5 {
			index := strings.LastIndex(tmps[3], ".")
			if index == -1 {
				continue
			}
			if port, err := strconv.Atoi(tmps[3][index+1:]); err == nil {
				p.NetStatUDP[port] = struct{}{}
			}
		}
	}
	return nil
}

func (p *PortCheck) runNetStatOne(proto int, port int) bool {
	var cmd string
	if proto == UDPProto {
		cmd = fmt.Sprintf("netstat -an | grep \"^%s\" | grep \"[^0-9]%d \" > /dev/null 2>&1", "udp", port)
	} else if proto == TCPProto {
		cmd = fmt.Sprintf("netstat -an | grep \"^%s\" | grep \"[^0-9]%d \" > /dev/null 2>&1", "tcp", port)
	} else {
		return false
	}
	log.Println(cmd)

	msg, ret := wolexec.ExecUnixCmd(cmd, 60)
	if ret != 0 {
		return false
	} else if msg == "" {
		return true
	}

	return false
}

func (p *PortCheck) connPort(proto int, port int) bool {
	var fd int
	var err error

	// 检查ipv4
	if proto == UDPProto {
		if fd, err = syscall.Socket(syscall.AF_INET, syscall.SOCK_DGRAM, syscall.IPPROTO_UDP); err != nil {
			return false
		}
	} else if proto == TCPProto {
		if fd, err = syscall.Socket(syscall.AF_INET, syscall.SOCK_STREAM, syscall.IPPROTO_TCP); err != nil {
			return false
		}
	} else {
		return false
	}
	sa := &syscall.SockaddrInet4{Port: port}
	copy(sa.Addr[:], net.IPv4zero.To4())
	if err := syscall.Bind(fd, sa); err != nil {
		syscall.Close(fd)
		return true
	}
	syscall.Close(fd)

	// 检查ipv6
	if proto == UDPProto {
		if fd, err = syscall.Socket(syscall.AF_INET6, syscall.SOCK_DGRAM, syscall.IPPROTO_UDP); err != nil {
			return false
		}
	} else if proto == TCPProto {
		if fd, err = syscall.Socket(syscall.AF_INET6, syscall.SOCK_STREAM, syscall.IPPROTO_TCP); err != nil {
			return false
		}
	} else {
		return false
	}
	sa6 := &syscall.SockaddrInet6{Port: port, ZoneId: uint32(0)}
	copy(sa6.Addr[:], net.IPv6zero.To16())
	if err := syscall.Bind(fd, sa6); err != nil {
		syscall.Close(fd)
		return true
	}
	syscall.Close(fd)

	return false
}

func (p *PortCheck) CheckPorts() error {
	if err := p.runNetStatAll(); err != nil {
		return fmt.Errorf("run netstat all error: %s", err.Error())
	}

	// 检查tcp
	for i := 0; i <= 65535; i++ {
		if p.connPort(TCPProto, i) {
			if _, exist := p.NetStatTCP[i]; exist {
				continue
			}

			// sleep 避免造成冲击
			time.Sleep(2 * time.Second)

			if !p.runNetStatOne(TCPProto, i) && p.connPort(TCPProto, i) {
				// TODO: 记录可能存在告警
				log.Println(fmt.Sprintf("Kernel-level rootkit or trojaned: %d", i))
			}

			// TODO: 告警数操作一定数据可提前截止
		}
	}
	// 检查udp
	for i := 0; i <= 65535; i++ {
		if p.connPort(UDPProto, i) {
			if _, exist := p.NetStatUDP[i]; exist {
				continue
			}

			// sleep 避免造成冲击
			time.Sleep(2 * time.Second)

			if !p.runNetStatOne(UDPProto, i) && p.connPort(UDPProto, i) {
				// TODO: 记录可能存在告警
				log.Println(fmt.Sprintf("Kernel-level rootkit or trojaned: %d", i))
			}

			// TODO: 告警数操作一定数据可提前截止
		}
	}

	return nil
}

func main2()  {
	p := &PortCheck{
		NetStatTCP: make(map[int]struct{}),
		NetStatUDP: make(map[int]struct{}),
	}
	if err := p.CheckPorts(); err != nil {
		log.Println(err.Error())
		return
	}

	return
}